#include "tree.hh"
#include "constants.hh"
#include "cppn/cppn.hh"
#include <list>
#include <iostream>
#include <memory>
#include <cassert>
using namespace std;

namespace brain {

    struct Tree::Leaf {
        float x, y, width;
        int level;


        Leaf (float x, float y, float width, int level)
            : x(x), y(y), width(width), level(level) {
            children.reserve(4);
        }

        vector<Leaf> children;
        float w;
        float variance = -1;
    };

    set<Connection> const&
    Tree::get_connections_for (Node const& node, bool outgoing) {

        if (outgoing && outgoing_connections.count(node) == 1)
            return outgoing_connections.find(node)->second;
        else if (!outgoing && incoming_connections.count(node) == 1)
            return incoming_connections.find(node)->second;

        unique_ptr<Leaf> root = make_unique<Leaf>(0, 0, 1, 0);

        list<Leaf*> queue;
        queue.push_back(root.get());

        while (!queue.empty()) {
            Leaf* p = queue.front(); queue.pop_front();

            p->children.push_back({p->x-p->width/2, p->y-p->width/2, p->width/2, p->level+1});
            p->children.push_back({p->x-p->width/2, p->y+p->width/2, p->width/2, p->level+1});
            p->children.push_back({p->x+p->width/2, p->y+p->width/2, p->width/2, p->level+1}); 
            p->children.push_back({p->x+p->width/2, p->y-p->width/2, p->width/2, p->level+1});

            for (int c = 0; c < 4; ++c) {
                Leaf* child = &p->children[c];

                if (outgoing) {
                    child->w = get_weight(node, child->x, child->y);
                } else {
                    child->w = get_weight(child->x, child->y, node);
                }
                
            }

            if (p->level < TREE_INITIAL_RESOLUTION
                || (p->level < TREE_MAX_RESOLUTION
                    && variance(p) > TREE_DIVISION_THRESHOLD)) 
            {
                for (int c = 0; c < 4; ++c) {
                    queue.push_back(&p->children[c]);
                }
            }

        }    

        set<Connection>& connections = outgoing ? outgoing_connections[node] : incoming_connections[node];
        prune(node, root.get(), outgoing, connections);
        return connections;
    }

    void
    Tree::prune (Node const& node, Leaf* leaf, bool outgoing, 
                 set<Connection>& connections) {
        
        for (int c = 0; c < 4; ++c) {
            if (leaf->children.size() == 0) return;
            Leaf* child = &leaf->children[c];

            if (variance(child) >= TREE_VARIANCE_THRESHOLD) {
                prune(node, child, outgoing, connections);

            } else {

                const float c_x = child->x;
                const float c_y = child->y;
                const float p_width = leaf->width;
                const float c_w = child->w;

                float dleft, dright, dtop, dbottom;

                if (c_w != 0) {
                    if (!outgoing) {
                        dleft = abs(c_w - get_weight(node, c_x-p_width, c_y)) / 2;
                        dright = abs(c_w - get_weight(node, c_x+p_width, c_y)) / 2;
                        dtop = abs(c_w - get_weight(node, c_x, c_y-p_width)) / 2;
                        dbottom = abs(c_w - get_weight(node, c_x, c_y+p_width)) / 2;
                    } else {
                        dleft = abs(c_w - get_weight(c_x-p_width, c_y, node)) / 2;
                        dright = abs(c_w - get_weight(c_x+p_width, c_y, node)) / 2;
                        dtop = abs(c_w - get_weight(c_x, c_y-p_width, node)) / 2;
                        dbottom = abs(c_w - get_weight(c_x, c_y+p_width, node)) / 2;
                    }

                    const float band = max(min(dtop, dbottom), min(dleft, dright));
                    if (band > TREE_BAND_THRESHOLD) {
    
                        if (outgoing) {
                          /*  cout << "new connection from (" << a << ", " << b << ") to ("
                                 << c_x << ", " << c_y << ") weight: " << c_w*WEIGHT_RANGE << "\n";*/
                            if (!connections.insert({node.get_x(), node.get_y(), c_x, c_y, 
                                    c_w*WEIGHT_ABS_RANGE}).second)
                               { }// cout << "\tconnection already exists!\n";
                        } else {
                          /*  cout << "new connection from (" << c_x << ", " << c_y << ") to ("
                                 << a << ", " << b << ") weight: " << c_w*WEIGHT_RANGE << endl;*/
                            if (!connections.insert({c_x, c_y, node.get_x(), node.get_y(),
                                    c_w*WEIGHT_ABS_RANGE}).second)
                               { }// cout << "\tconnection already exists!\n";
                        }
                    }
                }
            }
        }
    }

    void
    Tree::get_leaf_values (Leaf* leaf, std::vector<float>& list) {
        if (leaf->children.size() != 0) {//leaf != nullptr && leaf->children[0] != nullptr) {
            for (int c = 0; c < 4; ++c) {
                get_leaf_values(&leaf->children[c], list);
            }
        } else  {
            list.push_back(leaf->w);
        }
    }

    float
    Tree::variance (Leaf* leaf) {
        if (leaf->variance < 0) {
            if (leaf->children.size() == 0) {
                leaf->variance = 0;
                return 0;
            }
            
            vector<float> tmp;
            get_leaf_values(leaf, tmp);

            float mean = 0;
            for (float f: tmp) {
                mean += f;
            }
            mean /= tmp.size();

            float variance = 0;
            for (float f: tmp) {
                variance += (f - mean)*(f - mean);
            }
            variance /= tmp.size();
            leaf->variance = variance;
        }

        return leaf->variance;
    }

    float
    Tree::get_weight (float x1, float y1, float x2, float y2) {
        cppn->reset();
        cppn->set_input(x1, y1, x2, y2);
        float w = cppn->get_weight();

        return cppn->get_leo() ? w : 0;
    }

    float
    Tree::get_weight (Node const& src, Node const& dst) {
        return get_weight(src.get_x(), src.get_y(), dst.get_x(), dst.get_y());
    }

    float
    Tree::get_weight (Node const& src, float x, float y) {
        return get_weight(src.get_x(), src.get_y(), x, y);
    }

    float
    Tree::get_weight (float x, float y, Node const& dst) {
        return get_weight(x, y, dst.get_x(), dst.get_y());
    }

}











